"""
metrics_manager.py
==================

该模块实现了 MetricsManager —— 一个通用的指标管理器，用于在模型训练或验证过程中，
根据配置文件（MetricsConfig）动态构建和计算不同类型的评估指标。

支持的指标类型：
----------------
📊 分类指标：
    - accuracy
    - precision
    - recall
    - f1
    - auc

📈 回归指标：
    - mse
    - mae
    - r2

🧩 分割指标：
    - dice
    - iou

🧠 自定义指标：
    用户可通过 MetricsConfig.custom_metrics 传入字典形式的自定义函数。

使用方式：
----------
config = MetricsConfig(metrics=["accuracy", "f1", "mse"])
manager = MetricsManager(config)
results = manager.compute(preds, targets)
print(results)  # {'accuracy': 0.95, 'f1': 0.94, 'mse': 0.0023}
"""

import torch
from sklearn.metrics import (
    accuracy_score, precision_score, recall_score, f1_score, roc_auc_score,
    mean_squared_error, mean_absolute_error, r2_score
)
from typing import Dict, Callable
from config_manager import MetricsConfig


class MetricsManager:
    """
    统一的指标管理类，用于根据配置动态构造和计算模型评估指标。

    Attributes
    ----------
    config : MetricsConfig
        指标配置对象，定义使用哪些指标以及是否包含自定义指标。
    metric_fns : Dict[str, Callable]
        已注册的指标名称与对应计算函数的映射表。
    """

    def __init__(self, config: MetricsConfig):
        """
        初始化指标管理器。

        参数
        ----
        config : MetricsConfig
            包含指标类型和自定义指标的配置对象。
        """
        self.config = config
        self.metric_fns = self._build_metrics()

    def _build_metrics(self) -> Dict[str, Callable]:
        """
        根据配置文件动态构造指标函数映射表。

        返回
        ----
        Dict[str, Callable]
            指标名称 -> 指标计算函数
        """
        metrics = {}

        for name in self.config.metrics:
            name = name.lower()

            # === 分类指标 ===
            if name == "accuracy":
                metrics["accuracy"] = self._accuracy
            elif name == "precision":
                metrics["precision"] = self._precision
            elif name == "recall":
                metrics["recall"] = self._recall
            elif name == "f1":
                metrics["f1"] = self._f1
            elif name == "auc":
                metrics["auc"] = self._auc

            # === 回归指标 ===
            elif name == "mse":
                metrics["mse"] = self._mse
            elif name == "mae":
                metrics["mae"] = self._mae
            elif name == "r2":
                metrics["r2"] = self._r2

            # === 分割指标 ===
            elif name == "dice":
                metrics["dice"] = self._dice
            elif name == "iou":
                metrics["iou"] = self._iou

            else:
                raise ValueError(f"❌ 不支持的指标类型: {name}")

        # === 用户自定义指标 ===
        if self.config.custom_metrics:
            metrics.update(self.config.custom_metrics)

        return metrics

    def compute_metrics(self, preds, targets):
        """
        批量计算所有配置的指标。

        参数
        ----
        preds : torch.Tensor 或 np.ndarray
            模型预测结果。
        targets : torch.Tensor 或 np.ndarray
            对应的真实标签。

        返回
        ----
        Dict[str, float]
            每个指标名称及对应数值。
        """
        results = {}
        for name, fn in self.metric_fns.items():
            try:
                results[name] = float(fn(preds, targets))
            except Exception as e:
                results[name] = f"Error: {e}"
        return results

    # === 辅助方法 ===
    def _to_numpy(self, x):
        """将 torch.Tensor 安全地转换为 numpy 数组。"""
        if torch.is_tensor(x):
            return x.detach().cpu().numpy()
        return x

    # === 分类指标 ===
    def _accuracy(self, preds, targets):
        """分类准确率 (Accuracy)"""
        preds = self._to_numpy(preds)
        targets = self._to_numpy(targets)
        if preds.ndim > 1:
            preds = preds.argmax(axis=1)
        return accuracy_score(targets, preds)

    def _precision(self, preds, targets):
        """宏平均精确率 (Precision)"""
        preds = self._to_numpy(preds)
        targets = self._to_numpy(targets)
        if preds.ndim > 1:
            preds = preds.argmax(axis=1)
        return precision_score(targets, preds, average="macro", zero_division=0)

    def _recall(self, preds, targets):
        """宏平均召回率 (Recall)"""
        preds = self._to_numpy(preds)
        targets = self._to_numpy(targets)
        if preds.ndim > 1:
            preds = preds.argmax(axis=1)
        return recall_score(targets, preds, average="macro", zero_division=0)

    def _f1(self, preds, targets):
        """宏平均 F1 分数"""
        preds = self._to_numpy(preds)
        targets = self._to_numpy(targets)
        if preds.ndim > 1:
            preds = preds.argmax(axis=1)
        return f1_score(targets, preds, average="macro", zero_division=0)

    def _auc(self, preds, targets):
        """ROC AUC 指标（仅支持二分类）"""
        preds = self._to_numpy(preds)
        targets = self._to_numpy(targets)
        if preds.ndim > 1 and preds.shape[1] == 2:
            preds = preds[:, 1]
        return roc_auc_score(targets, preds)

    # === 回归指标 ===
    def _mse(self, preds, targets):
        """均方误差 (Mean Squared Error)"""
        preds = self._to_numpy(preds)
        targets = self._to_numpy(targets)
        return mean_squared_error(targets, preds)

    def _mae(self, preds, targets):
        """平均绝对误差 (Mean Absolute Error)"""
        preds = self._to_numpy(preds)
        targets = self._to_numpy(targets)
        return mean_absolute_error(targets, preds)

    def _r2(self, preds, targets):
        """R² 决定系数 (Coefficient of Determination)"""
        preds = self._to_numpy(preds)
        targets = self._to_numpy(targets)
        return r2_score(targets, preds)

    # === 分割指标 ===
    def _dice(self, preds, targets, eps=1e-7):
        """Dice 系数，用于语义分割任务"""
        if torch.is_tensor(preds):
            preds = torch.sigmoid(preds) if preds.dtype == torch.float32 else preds
            preds = (preds > 0.5).float()
            targets = targets.float()
            intersection = (preds * targets).sum()
            return (2. * intersection + eps) / (preds.sum() + targets.sum() + eps)
        else:
            preds = (preds > 0.5).astype(float)
            targets = targets.astype(float)
            intersection = (preds * targets).sum()
            return (2. * intersection + eps) / (preds.sum() + targets.sum() + eps)

    def _iou(self, preds, targets, eps=1e-7):
        """IoU（交并比），用于分割任务"""
        if torch.is_tensor(preds):
            preds = torch.sigmoid(preds) if preds.dtype == torch.float32 else preds
            preds = (preds > 0.5).float()
            targets = targets.float()
            intersection = (preds * targets).sum()
            union = preds.sum() + targets.sum() - intersection
            return (intersection + eps) / (union + eps)
        else:
            preds = (preds > 0.5).astype(float)
            targets = targets.astype(float)
            intersection = (preds * targets).sum()
            union = preds.sum() + targets.sum() - intersection
            return (intersection + eps) / (union + eps)

    def get_metrics(self):
        """
        返回所有已注册的指标函数。

        返回
        ----
        Dict[str, Callable]
            指标名称与对应函数。
        """
        return self.metric_fns
